import numpy as np
import pandas as pd
import nltk
import re
from collections import Counter,OrderedDict
import tensorflow.contrib.keras as kr
import os

def clean_str(string):
    """
    Tokenization/string cleaning for all datasets except for SST.
    Original taken from https://github.com/yoonkim/CNN_sentence/blob/master/process_data.py
    """
    string = re.sub(r"[^A-Za-z0-9(),!?\'\`]", " ", string)
    string = re.sub(r"\'s", " \'s", string)
    string = re.sub(r"\'ve", " \'ve", string)
    string = re.sub(r"n\'t", " n\'t", string)
    string = re.sub(r"\'re", " \'re", string)
    string = re.sub(r"\'d", " \'d", string)
    string = re.sub(r"\'ll", " \'ll", string)
    string = re.sub(r",", " , ", string)
    string = re.sub(r"!", " ! ", string)
    string = re.sub(r"\(", " \( ", string)
    string = re.sub(r"\)", " \) ", string)
    string = re.sub(r"\?", " \? ", string)
    string = re.sub(r"\s{2,}", " ", string)
    return string.strip().lower()


def load_data_and_labels(positive_data_file, negative_data_file):
    """
    Loads MR polarity data from files, splits the data into words and generates labels.
    Returns split sentences and labels.
    """
    # Load data from files
    positive_examples = list(open(positive_data_file, "r", encoding="UTF-8").readlines())
    positive_examples = [s.strip() for s in positive_examples]
    negative_examples = list(open(negative_data_file, "r", encoding="UTF-8").readlines())
    negative_examples = [s.strip() for s in negative_examples]
    # Split by words
    x_text = positive_examples + negative_examples
    x_text = [clean_str(sent) for sent in x_text]
    # Generate labels
    positive_labels = [[0, 1] for _ in positive_examples]
    negative_labels = [[1, 0] for _ in negative_examples]
    y = np.concatenate([positive_labels, negative_labels], 0)
    return [x_text, y]


def batch_iter(data, batch_size, num_epochs, shuffle=True):
    """
    Generates a batch iterator for a dataset.
    """
    data = np.array(data)
    data_size = len(data)
    num_batches_per_epoch = int((len(data) - 1) / batch_size) + 1
    for epoch in range(num_epochs):
        # Shuffle the data at each epoch
        if shuffle:
            shuffle_indices = np.random.permutation(np.arange(data_size))
            shuffled_data = data[shuffle_indices]
        else:
            shuffled_data = data
        for batch_num in range(num_batches_per_epoch):
            start_index = batch_num * batch_size
            end_index = min((batch_num + 1) * batch_size, data_size)
            yield shuffled_data[start_index:end_index]

def read_file(filename):
    """读取文件数据"""
    contents, labels = [], []
    with open(filename) as f:
        for line in f.readlines():
            try:
                datas = line.strip().split('\t')
                label_list = []
                content = None
                for i in range(len(datas)):
                    if i!=len(datas)-1:
                        label_list.append(datas[i])
                    else:
                        content = datas[i]#.replace(' ','')
                if content:
                    #contents.append(list(native_content(content)))
                    contents.append(content.split(' '))
                    labels.append(label_list)
            except:
                pass
    return contents, labels

def build_vocab(train_dir, vocab_dir, vocab_size=5000):
    """根据训练集构建词汇表，存储"""
    data_train, _ = read_file(train_dir)

    all_data = []
    for content in data_train:
        all_data.extend(content)

    counter = Counter(all_data)
    count_pairs = counter.most_common(vocab_size - 1)
    words, _ = list(zip(*count_pairs))
    # 添加一个 <PAD> 来将所有文本pad为同一长度
    words =['<PAD>']+list(words)
    with open(vocab_dir, 'w') as f:
        f.write('\n'.join(words) + '\n')
def read_vocab(vocab_dir):
    """读取词汇表"""
    with open(vocab_dir) as fp:
        words = [_.strip()for _ in fp.readlines()]
    word_to_id = dict(zip(words, range(len(words))))
    return words, word_to_id

# 获取标签字典
def get_labels_dict(path):
    label_list = []
    with open(path,encoding='utf-8') as file:
        contents = file.readlines()
        for content in contents:
            label = content.replace('\n','')
            if label not in label_list:
                label_list.append(label)
    labels = OrderedDict()
    count = 0
    for label in label_list:
        count+=1
        labels[label] = count
    return labels
current = os.curdir
def read_category(class_tag_dir):
    """读取分类目录，固定"""
    categories = get_labels_dict(current+class_tag_dir)
    categories = [x for x in categories.keys()]
    cat_to_id = dict(zip(categories, range(len(categories))))
    return categories, cat_to_id

def process_file(filename, word_to_id, cat_to_id, max_length=600):
    """将文件转换为id表示"""
    contents, labels = read_file(filename)

    data_id = []
    # 构造标签矩阵
    y_pad = np.zeros([len(contents), len(cat_to_id)])
    for i in range(len(contents)):
        data_id.append([word_to_id[x] for x in contents[i] if x in word_to_id])
        label_ids = [cat_to_id[cat] for cat in labels[i] if cat in cat_to_id]
        for label_id in label_ids:
            y_pad[i][int(label_id)]=1
    # 使用keras提供的pad_sequences来将文本pad为固定长度
    x_pad = kr.preprocessing.sequence.pad_sequences(data_id, max_length,padding='post',truncating='post')
    return x_pad, y_pad

# 获取准确率和召回率
def cal_rec_and_acc(predicted_labels, labels):
    label_no_zero = []
    for index, label in enumerate(labels):
        if int(label) == 1:
            label_no_zero.append(index)
    count = 0
    for predicted_label in predicted_labels:
        if int(predicted_label) in label_no_zero:
            count += 1
    rec = count / len(label_no_zero)
    acc = count / len(predicted_labels)
    return rec, acc
# 获取预测标签数组
def get_label_using_logits(logits, top_number=1):
    logits = np.ndarray.tolist(logits)
    predicted_labels = []
    for item in logits:
        index_list = np.argsort(item)[-top_number:]
        index_list = index_list[::-1]
        predicted_labels.append(np.ndarray.tolist(index_list))
    return predicted_labels





if __name__ == "__main__":
    pos_dir = "data/rt-polaritydata/rt-polarity.pos"
    neg_dir = "data/rt-polaritydata/rt-polarity.neg"

    load_data_and_labels(pos_dir, neg_dir)
